import pandas as pd
import numpy as np
import statsmodels.api as sm
from statsmodels.tsa.stattools import adfuller
import matplotlib.pyplot as plt

# 生成示例数据
np.random.seed(42)
dates = pd.date_range('2025-01-01', periods=100, freq='D')
Y = np.random.randn(100).cumsum()
data = pd.DataFrame({'Date': dates, 'Y': Y})
data.set_index('Date', inplace=True)

# 查看数据
print(data.head())


def generate_lags(df, column, max_lag):
    """
    生成指定列的滞后值
    :param df: DataFrame
    :param column: 需要生成滞后值的列名
    :param max_lag: 最大滞后阶数
    :return: 包含滞后值的DataFrame
    """
    for i in range(1, max_lag + 1):
        df[f'{column}_lag{i}'] = df[column].shift(i)
    return df


# 生成滞后值
max_lag = 2
data_with_lags = generate_lags(data, 'Y', max_lag)

# 删除含有NA值的行
data_with_lags.dropna(inplace=True)

# 查看生成滞后值后的数据
print(data_with_lags.head())


def check_stationarity(series):
    result = adfuller(series)
    print(f'ADF Statistic: {result[0]}')
    print(f'p-value: {result[1]}')
    for key, value in result[4].items():
        print(f'Critical Value ({key}): {value}')


check_stationarity(data_with_lags['Y'])
# 使用AIC准则选择最优滞后阶数
best_aic = float("inf")
best_order = 0

for p in range(1, 11):  # 尝试1到10阶滞后
    model = sm.tsa.ARIMA(data['Y'], order=(p, 0, 0))
    results = model.fit()
    if results.aic < best_aic:
        best_aic = results.aic
        best_order = p

print(f'Best lag order: {best_order}')

# 选择最优滞后阶数
p = best_order

# 构建并拟合自回归模型
X = data_with_lags[[f'Y_lag{i}' for i in range(1, p + 1)]]
X = sm.add_constant(X)
y = data_with_lags['Y']

model = sm.OLS(y, X)
results = model.fit()
print(results.summary())

# # 检查残差
# residuals = results.resid
#
# # 绘制残差的自相关图和偏自相关图
# fig, ax = plt.subplots(1, 2, figsize=(12, 5))
# sm.graphics.tsa.plot_acf(residuals, lags=20, ax=ax[0])
# sm.graphics.tsa.plot_pacf(residuals, lags=20, ax=ax[1])
# plt.show()
#
# # 检查残差的正态性
# import scipy.stats as stats
#
# # 绘制残差的QQ图
# stats.probplot(residuals, dist="norm", plot=plt)
# plt.title("Q-Q Plot of Residuals")
# plt.show()
# print(data.index[-1])
# # 预测未来10个时间点
future_dates = pd.date_range(start=data.index[-1],inclusive="right",periods=11,freq='D')
future_data = pd.DataFrame(index=future_dates, columns=data.columns)

# 生成未来数据的滞后值
for i in range(1, p + 1):
    ss = pd.Series([np.nan] * (10 - i))
    future_data[f'Y_lag{i}'] = pd.concat([data_with_lags[f'Y_lag{i}'].iloc[-i:],ss])
    print(future_data)
# future_data.dropna(inplace=True)
print(future_data)
# 添加常数项
future_X = sm.add_constant(future_data[[f'Y_lag{i}' for i in range(1, p + 1)]])
print(future_X)
# 进行预测
predictions = results.predict(future_X)
print(predictions)
# 绘制预测结果
plt.figure(figsize=(12, 6))
plt.plot(data.index, data['Y'], label='Observed')
plt.plot(future_data.index, predictions, color='r', label='Forecast')
plt.legend()
plt.show()